package ai.example.langchain4j.controller;

import ai.example.langchain4j.service.LangChain4jService;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import lombok.RequiredArgsConstructor;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.List;
import java.util.stream.Collectors;

@RestController
@RequestMapping("/api/rag")
@RequiredArgsConstructor
public class RagController {
  private final EmbeddingModel embeddingModel;
  private final LangChain4jService langChain4jService;

  @GetMapping("/vector")
  public String toVector(@RequestParam String message) {
    List<Float> vector = embeddingModel.embed(message).content().vectorAsList();
    return vector.stream().limit(10).toList().toString() + " … (total " + vector.size() + " dims)";
  }

  @PostMapping("/file")
  public String embedFileContent(@RequestParam(value = "file", required = false) MultipartFile file) throws IOException {
    return langChain4jService.upload(file);
  }

  @GetMapping("/related")
  public List<String> relatedQuery(@RequestParam String message) {
    List<EmbeddingMatch<TextSegment>> matches = langChain4jService.search(message);
    return matches.stream()
      .map(m -> String.format("Score: %.3f, Text: %s",
        m.score(), m.embedded().text()))
      .collect(Collectors.toList());
  }
}
